#include "data_utils.h"
#include <cstring>
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
extern void reflection_pad1d_do(uint32_t blockDim, 
                                  void *stream, 
                                  uint8_t *x, 
                                  uint8_t *paddings, 
                                  uint8_t *y);
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void reflection_pad1d(GM_ADDR x, GM_ADDR paddings, GM_ADDR y);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;
    size_t inputByteSize = 8 * 2 * 1024 * sizeof(float);
    size_t paddingsByteSize = 2 * sizeof(int32_t);

    int32_t pad_left = 0, pad_right = 0;
    uint8_t* paddingsTemp = (uint8_t*)malloc(paddingsByteSize);
    ReadFile("./input/paddings.bin", paddingsByteSize, paddingsTemp, paddingsByteSize);
    memcpy(&pad_left, paddingsTemp, sizeof(int32_t));
    memcpy(&pad_right, paddingsTemp + sizeof(int32_t), sizeof(int32_t));
    free(paddingsTemp); // 释放临时缓冲区

    size_t outputByteSize = 8 * 2 * (1024 + pad_left + pad_right) * sizeof(float);

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *x = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *paddings = (uint8_t *)AscendC::GmAlloc(paddingsByteSize);
    uint8_t *y = (uint8_t *)AscendC::GmAlloc(outputByteSize);

    ReadFile("./input/input_x.bin", inputByteSize, x, inputByteSize);
    ReadFile("./input/paddings.bin", paddingsByteSize, paddings, paddingsByteSize);

    AscendC::SetKernelMode(KernelMode::AIV_MODE);
    ICPU_RUN_KF(reflection_pad1d, blockDim, x, paddings, y);
    WriteFile("./output/output_y.bin", y, outputByteSize);

    AscendC::GmFree((void *)x);
    AscendC::GmFree((void *)paddings);
    AscendC::GmFree((void *)y);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));
    uint8_t *xHost, *paddingsHost, *yHost;
    uint8_t *xDevice, *paddingsDevice, *yDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&xHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&paddingsHost), paddingsByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&yHost), outputByteSize)); // 按正确大小分配
    CHECK_ACL(aclrtMalloc((void **)&xDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&paddingsDevice, paddingsByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&yDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST)); // 按正确大小分配

    ReadFile("./input/input_x.bin", inputByteSize, xHost, inputByteSize);
    ReadFile("./input/paddings.bin", paddingsByteSize, paddingsHost, paddingsByteSize);

    CHECK_ACL(aclrtMemcpy(xDevice, inputByteSize, xHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(paddingsDevice, paddingsByteSize, paddingsHost, paddingsByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    reflection_pad1d_do(blockDim, stream, xDevice, paddingsDevice, yDevice);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(yHost, outputByteSize, yDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_y.bin", yHost, outputByteSize);
    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(paddingsDevice));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFreeHost(xHost));
    CHECK_ACL(aclrtFreeHost(paddingsHost));
    CHECK_ACL(aclrtFreeHost(yHost));
    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}